package nsfs

import (
	"GLibHac/crypto"
	"GLibHac/tools"
	"encoding/binary"
	"github.com/pkg/errors"
)

const (
	Aes128XtsTransformBlockSize      = 128
	Aes128XtsTransformBlockSizeBytes = Aes128XtsTransformBlockSize / 8
)

type Aes128XtsTransform struct {
	key1       crypto.TransformFunc
	key2       crypto.TransformFunc
	decrypting bool
	cc         []byte
	pp         []byte
	t          []byte
}

func (t *Aes128XtsTransform) TransformBlock(buffer []byte, sector uint64) int {
	var lim int

	/* get number of blocks */
	count := len(buffer)
	var m = count >> 4
	var mo = count & 15
	var alignedCount = AlignUpIntUint(count, Aes128XtsTransformBlockSizeBytes)

	/* for i = 0 to m-2 do */
	if mo == 0 {
		lim = m
	} else {
		lim = m - 1
	}

	tweak := make([]byte, alignedCount)
	t.FillArrayFromSector(tweak, sector)

	/* encrypt the tweak */
	t.key2(tweak, tweak[:16])

	t.FillTweakBuffer(tweak[:alignedCount])
	if lim > 0 {
		XorArrays(buffer[:lim*16], tweak)
		t.key1(buffer, buffer[:lim*16])
		XorArrays(buffer[:lim*16], tweak)
	}

	if mo > 0 {
		copy(t.cc, tweak[lim*16:])

		if t.decrypting {
			copy(t.cc, tweak[lim*16+16:])

			/* CC = tweak encrypt block m-1 */
			t.TweakCrypt(buffer, t.pp, t.cc)

			/* Cm = first ptlen % 16 bytes of CC */
			var i int
			for i = 0; i < mo; i++ {
				t.cc[i] = buffer[16+i]
				buffer[16+i] = t.pp[i]
			}

			for ; i < 16; i++ {
				t.cc[i] = t.pp[i]
			}

			/* Cm-1 = Tweak encrypt PP */
			t.TweakCrypt(t.cc, buffer, t.t)
		} else {
			/* CC = tweak encrypt block m-1 */
			t.TweakCrypt(buffer, t.cc, t.t)

			/* Cm = first ptlen % 16 bytes of CC */
			var i int
			for i = 0; i < mo; i++ {
				t.pp[i] = buffer[16+i]
				buffer[16+i] = t.cc[i]
			}

			for ; i < 16; i++ {
				t.pp[i] = t.cc[i]
			}

			/* Cm-1 = Tweak encrypt PP */
			t.TweakCrypt(t.pp, buffer, t.t)
		}
	}

	return count
}

func (t *Aes128XtsTransform) FillArrayFromSector(value []byte, sector uint64) {
	for i := 0xF; i >= 0; i-- {
		value[i] = byte(sector)
		sector >>= 8
	}
}

func (t *Aes128XtsTransform) FillTweakBuffer(buf []byte) {
	a := binary.LittleEndian.Uint64(buf[8:])
	b := binary.LittleEndian.Uint64(buf)
	l := len(buf)
	for i := 16; i < l; i += 8 {
		tt := uint64(int64(a)>>63) & 0x87
		a = (a << 1) | (b >> 63)
		b = (b << 1) ^ tt
		binary.LittleEndian.PutUint64(buf[i:], b)
		i += 8
		binary.LittleEndian.PutUint64(buf[i:], a)
	}
}

func (t *Aes128XtsTransform) TweakCrypt(inputBuffer []byte, outputBuffer []byte, tt []byte) {
	for x := 0; x < 16; x++ {
		outputBuffer[x] = inputBuffer[x] ^ tt[x]
	}

	t.key1(outputBuffer[:16], outputBuffer)

	for x := 0; x < 16; x++ {
		outputBuffer[x] = outputBuffer[x] ^ tt[x]
	}
	t.MultiplyByX(tt)
}

func (t *Aes128XtsTransform) MultiplyByX(i []byte) {
	var t1 byte = 0
	var tt byte = 0

	for x := 0; x < 16; x++ {
		tt = i[x] >> 7
		i[x] = ((i[x] << 1) | t1) & 0xFF
		t1 = tt
	}

	if tt > 0 {
		i[0] ^= 0x87
	}
}

func NewAes128XtsTransform(key1, key2 []byte, decrypting bool) (this *Aes128XtsTransform, err error) {
	this = &Aes128XtsTransform{
		cc:         make([]byte, 16),
		pp:         make([]byte, 16),
		t:          make([]byte, 16),
		decrypting: decrypting,
	}
	if decrypting {
		this.key1, err = crypto.NewAesECBDecrypt(key1)
		if err != nil {
			return
		}
		this.key2, err = crypto.NewAesECBEncrypt(key2)
	} else {
		this.key1, err = crypto.NewAesECBEncrypt(key1)
		if err != nil {
			return
		}
		this.key2, err = crypto.NewAesECBEncrypt(key2)
	}
	return
}

const (
	Aes128CtrTransformBlockSize      = 128
	Aes128CtrTransformBlockSizeBytes = Aes128CtrTransformBlockSize / 8
)

type Aes128CtrTransform struct {
	Counter       []byte
	transformFunc crypto.TransformFunc
}

func (t *Aes128CtrTransform) TransformBlock(data []byte) {
	blockCount := tools.DivideUp(int64(len(data)), Aes128CtrTransformBlockSizeBytes)
	length := blockCount * Aes128CtrTransformBlockSizeBytes
	counterXor := make([]byte, length)
	copy(counterXor, t.Counter)
	FillDecryptedCounter(counterXor)
	t.transformFunc(counterXor, counterXor)
	XorArrays(data, counterXor)
}

func FillDecryptedCounter(buffer []byte) {
	hi := binary.LittleEndian.Uint64(buffer)
	lo := binary.BigEndian.Uint64(buffer[8:])
	for i := 16; i < len(buffer); i += 8 {
		lo += 1
		binary.LittleEndian.PutUint64(buffer[i:], hi)
		i += 8
		binary.BigEndian.PutUint64(buffer[i:], lo)
	}
}

func NewAes128CtrTransform(key, counter []byte) (r *Aes128CtrTransform, err error) {
	if key == nil {
		return nil, errors.Errorf("key cannot be null")
	}
	if counter == nil {
		return nil, errors.Errorf("counter cannot be null")
	}
	if len(key) != Aes128CtrTransformBlockSizeBytes {
		return nil, errors.Errorf("key must be %d bytes long", Aes128CtrTransformBlockSizeBytes)
	}
	if len(counter) != Aes128CtrTransformBlockSizeBytes {
		return nil, errors.Errorf("counter must be %d bytes long", Aes128CtrTransformBlockSizeBytes)
	}
	transformFunc, err := crypto.NewAesECBEncrypt(key)
	if err != nil {
		return nil, err
	}
	r = &Aes128CtrTransform{
		Counter:       make([]byte, Aes128CtrTransformBlockSizeBytes),
		transformFunc: transformFunc,
	}
	copy(r.Counter, counter[:Aes128CtrTransformBlockSizeBytes])
	return r, nil
}
